package com.neuedu.it.newidea.web;

import com.neuedu.it.newidea.po.User;
import com.neuedu.it.newidea.web.dto.RespEntity;
import jakarta.annotation.PostConstruct;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.InputStreamResource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/*
 * 学校活动AI控制器。
 * 提供管理员上传活动PDF并向量化、用户查询活动信息的功能。
*/
@RestController
@CrossOrigin
@RequestMapping("/ai/school-activities")
public class SchoolActivitiesController {

    @Autowired
    @Qualifier("chatWithSchoolActivitiesClient")
    private ChatClient chatClientSchoolActivities;

    @Autowired
    private EmbeddingModel embeddingModel;

    // 向量化JSON文件的存储目录
    @Value("${school_activities.vector-store-dir}")
    private String vectorStoreDir;

    // 存储不同学校活动文件对应的向量存储实例，以文件的basename作为key
    private final Map<String, VectorStore> fileVectorStores = new ConcurrentHashMap<>();

    // 获取向量存储JSON文件的基准路径
    private Path getVectorStoreBasePath() {
        // 使用项目根目录下的相对路径
        String projectRoot = System.getProperty("user.dir");
        return Paths.get(projectRoot, vectorStoreDir);
    }


    @PostConstruct
    public void init() {
        try {
            Path vectorStorePath = getVectorStoreBasePath();
            Files.createDirectories(vectorStorePath);
            // 加载所有已有的.json向量文件
            loadAllVectorStores();
            System.out.println("应用启动时自动加载所有已存在的学校活动向量存储文件到内存。");
        } catch (IOException e) {
            System.err.println("学校活动向量存储文件失败: " + e.getMessage());
        }
    }

    /**
     * 管理员上传学校活动PDF文件，并将其向量化存储。
     * **只有score为-1的用户（管理员）才能上传。**
     * **不保存原始PDF文件，只保存向量化后的JSON文件。**
     *
     * @param request HttpServletRequest 用于获取当前用户（JWT拦截器设置的）
     * @param file 待上传的PDF文件
     * @return 包含操作结果的RespEntity
     * @throws IOException 如果文件处理失败
     */
    @PostMapping("/upload")
    public RespEntity uploadSchoolActivityFile(HttpServletRequest request, @RequestParam("file") MultipartFile file) throws IOException {
        User curUser = (User) request.getAttribute("curUser");
        // 权限校验：检查用户是否为管理员 (score == -1)
        if(curUser != null) {
            System.out.println(curUser);
        }
        if (curUser == null || curUser.getScore() != -1) {
            return new RespEntity("4003", "只有管理员才能上传学校活动文件", null);
        }

        if (file.isEmpty()) {
            return new RespEntity("4000", "请选择要上传的PDF文件", null);
        }

        // 验证文件类型是否为PDF
        if (!"application/pdf".equalsIgnoreCase(file.getContentType())) {
            return new RespEntity("4000", "请上传PDF格式的文件", null);
        }

        String originalFileName = file.getOriginalFilename();
        if (originalFileName == null || originalFileName.trim().isEmpty()) {
            return new RespEntity("4000", "文件名无效", null);
        }

        // 从文件名中提取basename (不含扩展名)
        String baseName = originalFileName.contains(".") ?
                originalFileName.substring(0, originalFileName.lastIndexOf('.')) : originalFileName;

        // 构建向量存储JSON文件的完整路径
        Path vectorStoreFilePath = getVectorStoreBasePath().resolve(baseName + ".json");
        File vectorStoreFile = vectorStoreFilePath.toFile();

        // 检查是否已存在该文件的向量存储，如果存在则直接加载并返回
        if (vectorStoreFile.exists()) {
            try {
                SimpleVectorStore store = SimpleVectorStore.builder(embeddingModel).build();
                store.load(vectorStoreFile);
                fileVectorStores.put(baseName, store);
                System.out.println("已加载学校活动文件的向量文件: " + vectorStoreFilePath);
                return new RespEntity("2000", "该学校活动文件已存在，已加载现有向量存储", Map.of("fileName", originalFileName));
            } catch (Exception e) {
                System.err.println("Failed to load existing school activities vector store: " + e.getMessage());
                return new RespEntity("5000", "加载现有向量存储失败: " + e.getMessage(), null);
            }
        }

        // 直接从MultipartFile的InputStream处理PDF
        try (var inputStream = file.getInputStream()) {
            PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(
                    new InputStreamResource(inputStream),
                    PdfDocumentReaderConfig.builder()
                            .withPageTopMargin(0)
                            .withPageExtractedTextFormatter(( ExtractedTextFormatter.builder().withNumberOfTopTextLinesToDelete(0).build()))
                            .withPagesPerDocument(1)
                            .build()
            );

            List<Document> documentList = pdfReader.read();
            System.out.println("成功读取 " + documentList.size() + " document chunks from provided PDF.");

            List<Document> cleanedDocumentList = documentList.stream()
                    .map(doc -> {
                        Map<String, Object> originalMetadata = doc.getMetadata();
                        Map<String, Object> cleanedMetadata = originalMetadata.entrySet().stream()
                                .filter(e -> e.getKey() != null && e.getValue() != null)
                                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
                        return new Document(doc.getText(), cleanedMetadata);
                    })
                    .collect(Collectors.toList());

            SimpleVectorStore documentVectorStore = SimpleVectorStore.builder(embeddingModel).build();

            documentVectorStore.add(cleanedDocumentList);
            System.out.println("学校活动文件向量化成功");

            documentVectorStore.save(vectorStoreFile);
            fileVectorStores.put(baseName, documentVectorStore);
            System.out.println("学校活动向量化文件保存至: " + vectorStoreFilePath);

        } catch (Exception e) {
            e.printStackTrace();
            return new RespEntity("5000", "处理学校活动文件失败: " + e.getMessage(), null);
        }

        return new RespEntity("2000", "学校活动信息上传并向量化成功", Map.of("fileName", originalFileName));
    }

    /**
     * 用户查询学校活动信息。
     * AI将根据用户问题，在已上传的学校活动文档中进行检索（RAG模式）并作答。
     *
     * @param body 请求体，包含"question" (用户提问) 和可选的 "fileName" (指定要查询的文件名，不含扩展名)。
     *             如果未指定fileName，则会在所有已加载的向量存储中进行搜索。
     * @return AI的回答
     */
    @PostMapping("/query")
    public RespEntity querySchoolActivities(@RequestBody Map<String, String> body) {
        String question = body.get("question");
        // 可选参数：指定特定文件查询
        String fileName = body.get("fileName");

        if (question == null || question.trim().isEmpty()) {
            return new RespEntity("4000", "问题不能为空", null);
        }

        List<Document> relevantDocuments = new java.util.ArrayList<>();
        if (fileName != null && !fileName.trim().isEmpty()) {
            // 如果指定了文件名，则只在该文件的向量存储中搜索
            String baseName = fileName.contains(".") ?
                    fileName.substring(0, fileName.lastIndexOf('.')) : fileName;
            VectorStore specificVectorStore = fileVectorStores.get(baseName);
            if (specificVectorStore == null) {
                return new RespEntity("4000", "未找到指定学校活动文件的向量存储: " + fileName, null);
            }
            relevantDocuments.addAll(specificVectorStore.similaritySearch(
                    SearchRequest.builder().query(question).topK(5).similarityThreshold(0.1).build()
            ));
        } else {
            // 如果未指定文件名，则在所有已加载的学校活动向量存储中搜索
            if (fileVectorStores.isEmpty()) {
                return new RespEntity("4004", "当前没有可查询的学校活动文件，请联系管理员上传", null);
            }
            for (VectorStore vs : fileVectorStores.values()) {
                relevantDocuments.addAll(vs.similaritySearch(
                        SearchRequest.builder().query(question).topK(5).similarityThreshold(0.1).build()
                ));
            }
        }

        // 将检索到的相关文档内容格式化为上下文
        String context = relevantDocuments.stream()
                .map(document -> "内容:\n" + document.getText() + "\n元数据:" + document.getMetadata())
                .collect(Collectors.joining("\n---\n"));

        // 修改finalPrompt，要求AI返回特定格式
        String finalPrompt = "用户提问: " + question + "\n\n以下是学校活动相关的参考信息:\n" + context +
                "\n\n请根据提供的参考信息，以以下结构化格式回答：\n" +
                "活动标题：\n" +
                "- 服务时间：[时间信息]\n" +
                "- 服务地点：[地点信息]\n" +
                "- 招募人数：[人数信息]\n" +
                "- 报名方式：[报名信息]\n" +
                "- 备注：[备注信息]\n\n" +
                "其他活动：[如有其他活动请简要列出]";
        if (context.isEmpty()) {
            finalPrompt = finalPrompt + "用户提问: " + question + "\n\n提示: 在学校活动信息库中没有找到相关的文档信息。请尝试根据你自己的背景知识回答，如果无法回答请说明。";
        } else {
            finalPrompt = finalPrompt + "用户提问: " + question + "\n\n以下是学校活动相关的参考信息:\n" + context + "\n\n请根据提供的参考信息，以简洁明了的方式，回答用户关于学校活动的问题。";
        }

        try {
            // AI生成回答
            String aiResponseContent = chatClientSchoolActivities.prompt()
                    .system("你是一个专业的信息查询助手，根据提供的文本资料回答用户关于学校活动的问题。")
                    .user(finalPrompt)
                    .call()
                    .content();
            return new RespEntity("2000", "AI查询成功", aiResponseContent);
        } catch (Exception e) {
            e.printStackTrace();
            return new RespEntity("5000", "AI查询服务异常: " + e.getMessage(), null);
        }
    }

    /**
     * 获取已上传的学校活动文件列表，即已成功生成向量存储的列表
     * 返回文件名列表 (不含.json扩展名)。
     */
    @GetMapping("/files")
    public RespEntity getUploadedSchoolActivityFiles() {
        try {
            Path vectorStorePath = getVectorStoreBasePath();
            if (!Files.exists(vectorStorePath)) {
                return new RespEntity("2000", "没有找到学校活动文件存储目录", List.of());
            }
            List<String> files = Files.list(vectorStorePath)
                    .filter(path -> path.toString().endsWith(".json"))
                    .map(path -> path.getFileName().toString().replace(".json", ""))
                    .collect(Collectors.toList());
            return new RespEntity("2000", "获取学校活动文件列表成功", files);
        } catch (IOException e) {
            return new RespEntity("5000", "获取学校活动文件列表失败: " + e.getMessage(), null);
        }
    }

    //加载所有已存在的向量存储文件
    private void loadAllVectorStores() throws IOException {
        Path vectorStorePath = getVectorStoreBasePath();
        if (Files.exists(vectorStorePath)) {
            Files.list(vectorStorePath)
                    .filter(path -> path.toString().endsWith(".json"))
                    .forEach(path -> {
                        try {
                            String fileName = path.getFileName().toString().replace(".json", "");
                            SimpleVectorStore store = SimpleVectorStore.builder(embeddingModel).build();
                            store.load(path.toFile());
                            fileVectorStores.put(fileName, store);
                            System.out.println("加载学校活动向量存储文件: " + fileName);
                        } catch (Exception e) {
                            System.err.println("加载学校活动向量存储文件失败: " + path.getFileName() + "，错误信息: " + e.getMessage());
                        }
                    });
        }
    }
}